import smbus
import time
import math
import numpy as np

PWR_MGMT_1   = 0x6B
SMPLRT_DIV   = 0x19
CONFIG       = 0x1A
GYRO_CONFIG  = 0x1B
INT_ENABLE   = 0x38
ACCEL_XOUT_H = 0x3B
ACCEL_YOUT_H = 0x3D
ACCEL_ZOUT_H = 0x3F
GYRO_XOUT_H  = 0x43
GYRO_YOUT_H  = 0x45
GYRO_ZOUT_H  = 0x47

CalGx = -14.51
CalGy = -48.631
CalGz = 5.116
CalAx = -809.336
CalAy = -461.796
CalAz = 1203.156

Aunit = 16384.0
Gunit = 131.0


class Kalman():
    def __init__(self):
        self.P = np.matrix([[0., 0.],
                            [0., 0.]])
    def setKalmanAngle(self, angle):
        self.State = np.matrix([[angle],
                                [0.   ]])
    def getKalmanAngle(self, angle, gyro_rate, dt):
        R = 0.03
        Q = np.matrix([[0.001, 0.   ],
                       [0.,    0.003]])
        H = np.matrix( [1.,    0.   ])

        F = np.matrix([[1., -dt],
                       [0., 1. ]])
        B = np.matrix([[dt],
                       [0.]])
        #(I). State prediction
        self.State = F * self.State + B * gyro_rate
        #print('I. State prediction: \n', self.State)

        #(II). Covariance prediction
        self.P = F * self.P * np.transpose(F) + Q
        #print('II. Covariance prediction P: \n', self.P)

        #(III). Innovation
        I = angle - H * self.State
        #print('III. Innovation I: \n', I)

        #(IV). Innovation covariance S
        S = H * self.P * np.transpose(H) + R
        #print('IV. Innovation covariance S: \n', S)

        #(V). Kalman gain KG
        KG = self.P * np.transpose(H) / S
        #print('V. Kalman Gain: \n', KG)

        #(VI). Update state
        self.State = self.State + KG * I
        #print('VI. Update State: \n', self.State)

        #(VII). Update covariance
        self.P = (np.eye(2) - KG * H) * self.P
        #print('VII. Update Covariance P: \n', self.P)
        return self.State.item(0)


def MPU_Init():
    bus.write_byte_data(Device_Address, SMPLRT_DIV, 7)
    bus.write_byte_data(Device_Address, PWR_MGMT_1, 1)
    bus.write_byte_data(Device_Address, CONFIG, 0)
    bus.write_byte_data(Device_Address, GYRO_CONFIG, 24)
    bus.write_byte_data(Device_Address, INT_ENABLE, 1)


def read_raw_data(addr):
    high = bus.read_byte_data(Device_Address, addr)
    low = bus.read_byte_data(Device_Address, addr + 1)
    value = ((high << 8) | low)
    if (value >= 0x8000):
        return -((65535 - value) + 1)
    else:
        return value
    return value


def get_pitch(x, y, z): #y-rotation
    radians = math.atan2(x, math.sqrt(y*y + z*z))
    return math.degrees(radians)


def get_roll(x, y, z): #x-rotation
    radians = -math.atan2(y, z)
    return math.degrees(radians)


def get_yew(x, y, z):
    radians = -math.atan2(y, z)
    return math.degrees(radians)


bus = smbus.SMBus(1)
Device_Address = 0x68
MPU_Init()

print (" Reading Data of Gyroscope and Accelerometer")

# for i in range(Calt):
#     time.sleep(0.01)
#     CalAx += read_raw_data(ACCEL_XOUT_H)
#     CalAy += read_raw_data(ACCEL_YOUT_H)
#     CalAz += read_raw_data(ACCEL_ZOUT_H)
#     CalGx += read_raw_data(GYRO_XOUT_H)
#     CalGy += read_raw_data(GYRO_YOUT_H)
#     CalGz += read_raw_data(GYRO_ZOUT_H)
# CalAx /= Calt
# CalAy /= Calt
# CalAz /= Calt
# CalGx /= Calt
# CalGy /= Calt
# CalGz /= Calt
time_pre = time.time()
while True:
    time.sleep(1)
    acc_x = read_raw_data(ACCEL_XOUT_H) - CalAx
    acc_y = read_raw_data(ACCEL_YOUT_H) - CalAy
    acc_z = read_raw_data(ACCEL_ZOUT_H) - CalAz
    gyro_x = read_raw_data(GYRO_XOUT_H) - CalGx
    gyro_y = read_raw_data(GYRO_YOUT_H) - CalGy
    gyro_z = read_raw_data(GYRO_ZOUT_H) - CalGz
    Ax = acc_x/Aunit
    Ay = acc_y/Aunit
    Az = acc_z/Aunit
    Gx = gyro_x/Gunit
    Gy = gyro_y/Gunit
    Gz = gyro_z/Gunit
    dt = time.time() - time_pre
    R = np.sqrt(Ax ** 2 + Ay ** 2 + Az ** 2)
    
    print(Gx, Gy, Gz)